package org.example.ratelimit2.impl;

import org.example.ratelimit2.MyRateLimiter;

import java.time.LocalDateTime;
import java.time.ZoneOffset;
import java.util.Iterator;
import java.util.Map;
import java.util.Set;
import java.util.TreeMap;

/**
 * 通过滑动窗口，每秒限制10个
 * 1秒1个格子
 */
public class SlidingWindowRateLimiter2 implements MyRateLimiter {

    private final long permitsPerMinute;

    private final long timePerWindow;

    private final long limitTime;

    /**
     * 所有子窗口
     *  k：窗口时间
     *  v：该窗口的计数
     */
    private TreeMap<Long, Integer> counters;

    public SlidingWindowRateLimiter2(long permitsPerMinute, long timePerWindow, long limitTime) {
        this.permitsPerMinute = permitsPerMinute;
        this.timePerWindow = timePerWindow;
        this.limitTime = limitTime;
        counters = new TreeMap<>();
    }

    @Override
    public synchronized boolean tryAcquire() {
        // 假设每10秒一个窗口
        long currentWindowTime = LocalDateTime.now().toEpochSecond(ZoneOffset.UTC) / timePerWindow * timePerWindow;
        int count = getCurrentWindowCount(currentWindowTime);
        if (count >= permitsPerMinute) {
            return false;
        }
        // 当前窗口计数器+1
        counters.merge(currentWindowTime, 1, Integer::sum);
        return true;
    }

    public int getCurrentWindowCount(long currentTime) {
        long startTime = currentTime - limitTime + timePerWindow;
        int result = 0;
        Set<Map.Entry<Long, Integer>> entrySet = counters.entrySet();
        Iterator<Map.Entry<Long, Integer>> iterator = entrySet.iterator();
        while (iterator.hasNext()) {
            Map.Entry<Long, Integer> entry = iterator.next();
            if (entry.getKey() < startTime) {
                iterator.remove();
            } else {
                result = result + entry.getValue();
            }
        }
        return result;
    }
}
